package zhxt.helmet.service;
import zhxt.helmet.utils.FrameToMatConverter;

import org.bytedeco.javacpp.FloatPointer;
import org.bytedeco.opencv.global.opencv_core;
import org.bytedeco.opencv.global.opencv_imgproc;
import org.bytedeco.opencv.global.opencv_imgcodecs;
import org.bytedeco.opencv.opencv_core.*;
import ai.onnxruntime.*;
import org.springframework.stereotype.Service;

import java.nio.FloatBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class Yolov10Service {
    //直接返回符合格式要求的数据, 用于测试其他模块
    public List<Map<String, Object>> analogDetect() {
        // 定义包含数据的字符串
        String data = "2 0.913633 0.589057 0.086207 0.260727 0.915881\n" +
                "2 0.141715 0.639393 0.052257 0.176286 0.854267\n" +
                "1 0.157052 0.646877 0.014143 0.029273 0.657205\n" +
                "1 0.971764 0.593952 0.020381 0.038285 0.634854";

        // 创建 List<float[]> 类型的变量
        List<float[]> list = new ArrayList<>();

        // 按行分割数据
        String[] lines = data.split("\n");

        // 遍历每一行
        for (String line : lines) {
            // 按空格分割每行数据
            String[] values = line.split(" ");
            // 创建一个 float 数组来存储当前行的数据
            float[] floatArray = new float[values.length];
            // 将字符串转换为 float 类型并存储到数组中
            for (int i = 0; i < values.length; i++) {
                floatArray[i] = Float.parseFloat(values[i]);
            }
            // 将 float 数组添加到 List 中
            list.add(floatArray);
        }

        return getBoxes(list);
    }
    //检测图片
    public List<Map<String, Object>>  detect(Mat image) throws OrtException {
        // 1. 加载 ONNX 模型
        String modelPath = "D:\\codes\\java\\spring\\helmet\\best.onnx";
        OrtEnvironment env = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions sessionOptions = new OrtSession.SessionOptions();
        OrtSession session = env.createSession(modelPath, sessionOptions);

        // 2. 预处理图像
        int targetSize = 896;
        float confThreshold = 0.5f;

        // 2.1 保持长宽比缩放并填充
        Mat resized = resizeWithPadding(image, targetSize);
        //FrameToMatConverter.saveMatToFile(resized,"output_resized.jpg");//保存图片看看输入数据是否正常


        // 2.2 转换为 RGB
        Mat rgb = new Mat();
        opencv_imgproc.cvtColor(resized, rgb, opencv_imgproc.COLOR_BGR2RGB);
        resized.close();  // 释放原Mat
        //FrameToMatConverter.saveMatToFile(rgb,"output_rgb.jpg");//保存图片看看输入数据是否正常
        //归一化
        Mat floatMat = new Mat();
        rgb.convertTo(floatMat, opencv_core.CV_32F, 1.0 / 255.0, 0);
        rgb.close();  // 释放RGB Mat
        resized = floatMat;
        //FrameToMatConverter.saveMatToFile(resized,"output_float.jpg");//保存图片看看输入数据是否正常

        // 2.3 转换为 CHW 格式的 float 数组
        float[] inputData = matToCHW(resized, targetSize);

        // 3. 创建输入张量
        long[] inputShape = {1, 3, targetSize, targetSize};
        OnnxTensor inputTensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), inputShape);

        // 4. 运行推理
        Map<String, OnnxTensor> inputs = new HashMap<>();
        inputs.put("images", inputTensor); // 输入节点名称需与模型匹配
        OrtSession.Result results = session.run(inputs);

        // 5. 解析输出张量
        OnnxTensor outputTensor = (OnnxTensor) results.get("output0").get(); // 输出节点名称需与模型匹配
        float[][][] outputData = (float[][][]) outputTensor.getValue();

        // 6. 后处理：提取检测框并过滤置信度
        List<float[]> detections = parseOutputXYWHC(outputData[0], confThreshold,targetSize, image.cols(), image.rows());
        // 7. 转换输出格式方便网页显示
        List<Map<String, Object>> boxes = getBoxes(detections);
        // 7. 打印结果
//        for (float[] box : detections) {
//            System.out.printf("%d %f %f %f %f %f%n",
//                    (int) box[0], box[1], box[2], box[3], box[4], box[5]);
//        }
//        //直接输出
//        for(float[][] subArray : outputData){
//            System.out.println("subArray:");
//            for (float[] row : subArray) {
//                for (float value : row) {
//                    System.out.print(value + " ");
//                }
//                System.out.println();
//            }
//        }

        // 8. 释放资源
        inputTensor.close();
        outputTensor.close();
        session.close();
        env.close();
        resized.close();
        rgb.close();

        return boxes;
    }

    // 保持长宽比缩放并填充为正方形
    private Mat resizeWithPadding(Mat image, int targetSize) {
        double scale = Math.min((double) targetSize / image.cols(), (double) targetSize / image.rows());
        int newWidth = (int) (image.cols() * scale);
        int newHeight = (int) (image.rows() * scale);

        Mat resized = new Mat();
        opencv_imgproc.resize(image, resized, new Size(newWidth, newHeight));

        Mat padded = new Mat(targetSize, targetSize, opencv_core.CV_8UC3, new Scalar(114, 114, 114, 0));

        if (image.cols() > image.rows()) { // width > height
            Rect roi = new Rect(0, (targetSize - newHeight) / 2, newWidth, newHeight);
            Mat roiMat = padded.apply(roi);
            resized.copyTo(roiMat);
            roiMat.close();
        } else {
            Rect roi = new Rect((targetSize - newWidth) / 2, 0, newWidth, newHeight);
            Mat roiMat = padded.apply(roi);
            resized.copyTo(roiMat);
            roiMat.close();
        }

        resized.close();
        return padded;
    }

    // 将 OpenCV Mat 转换为 CHW 格式的 float 数组
    /**
     * 将 org.bytedeco.opencv 的 Mat 转换为 CHW 格式的 float 数组
     * @param mat 输入的Mat (CV_32FC3 类型，即float型的3通道矩阵)
     * @param targetSize 目标尺寸
     * @return CHW格式的float数组 (顺序: Channel-Height-Width)
     */
    private float[] matToCHW(Mat mat, int targetSize) {


        float[] data = new float[3 * targetSize * targetSize];
        int idx = 0;

        // 获取Mat的数据指针
        FloatPointer floatPtr = new FloatPointer(mat.data());
        long step = mat.step() / 4; // 步长(以float为单位)

        // 转换为CHW格式
        for (int c = 0; c < 3; c++) {          // Channel-first
            for (int y = 0; y < targetSize; y++) {     // Height
                for (int x = 0; x < targetSize; x++) {  // Width
                    // 计算内存位置: y * step + x * channels + c
                    data[idx++] = floatPtr.get(y * step + x * 3 + c);
                }
            }
        }

        return data;
    }

    //解析YOLOv10输出，将结果转为归一化的，相对于原图的：左上x 左上y w h confidence
    //原始输出的值除了置信度都不是归一化的，是在targetSize尺寸的图像上的绝对位置
    private List<float[]> parseOutputXYWHC(float[][] output, float confThreshold, int targetSize, int imgWidth, int imgHeight) {
        List<float[]> detections = new ArrayList<>();
        double scale = Math.min((double) targetSize / imgWidth, (double) targetSize / imgHeight);
        int newWidth = (int) (imgWidth * scale);
        int newHeight = (int) (imgHeight * scale);
        for (float[] prediction : output) {
            float confidence = prediction[4]; // 置信度在索引4
            if (confidence < confThreshold) continue;

            // 将targetSize = 896图上的xyxy框解析为原图的边界框（xywh 格式）
            float x,y, width, height;

            if(imgWidth>imgHeight) {
                x = prediction[0]/targetSize;
                y = (prediction[1]-(float) (targetSize - newHeight) /2)/newHeight;
                width = (prediction[2]-prediction[0])/targetSize;
                height =(prediction[3]-prediction[1])/newHeight;
            }
            else{
                x = (prediction[0]-(float) (targetSize - newWidth) /2)/newWidth;
                y = prediction[1]/targetSize;
                width = (prediction[2]-prediction[0])/newWidth;
                height =(prediction[3]-prediction[1])/targetSize;
            }

            detections.add(new float[]{prediction[5], x,y,width,height,confidence});
        }
        return detections;
    }

    // 解析 YOLOv10 输出（输出形状为 [1, 300, 6]），将结果转为labelimg能识别的格式（xCenter，yCenter，w,h,confidence）
    private List<float[]> parseOutputXcYcWHC(float[][] output, float confThreshold, int targetSize, int imgWidth, int imgHeight) {
        List<float[]> detections = new ArrayList<>();
        double scale = Math.min((double) targetSize / imgWidth, (double) targetSize / imgHeight);
        int newWidth = (int) (imgWidth * scale);
        int newHeight = (int) (imgHeight * scale);
        for (float[] prediction : output) {
            float confidence = prediction[4]; // 置信度在索引4
            if (confidence < confThreshold) continue;

            // 将targetSize = 896图上的xyxy框解析为原图的边界框（xywh 格式）
            float xCenter, yCenter, width, height;
            if(imgWidth>imgHeight) {
                xCenter = (prediction[0]+prediction[2])/2/targetSize;
                yCenter = ((prediction[1]+prediction[3])/2- (float) (targetSize - newHeight) /2)/newHeight;
                width = (prediction[2]-prediction[0])/targetSize;
                height =(prediction[3]-prediction[1])/newHeight;//框h/新图h
            }
            else{
                xCenter = ((prediction[0]+prediction[2])/2- (float) (targetSize - newWidth) /2)/newWidth;
                yCenter = (prediction[1]+prediction[3])/2/targetSize;
                width = (prediction[2]-prediction[0])/newWidth;
                height =(prediction[3]-prediction[1])/targetSize;
            }

            detections.add(new float[]{prediction[5], xCenter,yCenter,width,height,confidence});
        }
        return detections;
    }

    // 转换输出格式, 封装为boxes, 方便网页显示
    private List<Map<String, Object>> getBoxes(List<float[]> detections){
        List<Map<String, Object>> boxes = new ArrayList<>();
        for (float[] result : detections) {
            Map<String, Object> box = new HashMap<>();
            box.put("x", result[1]);       // 左上x
            box.put("y", result[2]);       // 左上y
            box.put("w", result[3]);   // 宽度
            box.put("h", result[4]);  // 高度
            box.put("c", result[5]); // 置信度
            if((int)result[0]==0) box.put("l", "h");    // 类别名称 <label, helmet>
            else if((int)result[0]==1) box.put("l", "w"); // <label, without>
            else box.put("l", "t"); //<label , two_wheels>
            boxes.add(box);
        }
        return boxes;
    }
}
